import torch

# torch.broadcast_tensors 是 PyTorch 中的一个函数，用于对一组张量进行广播操作，使它们在形状上变得兼容，以便后续的运算可以顺利进行。
# 广播是指将较小形状的张量扩展为与较大形状的张量一致的操作，而无需实际复制数据，从而高效地进行计算。
#
# 函数签名
# torch.broadcast_tensors(*tensors) -> List[Tensor]
# 参数
# *tensors：任意数量的输入张量。输入的每个张量的形状不一定相同，但必须满足广播规则。
# 返回值
# 返回一个张量列表，其中每个张量都被广播到了共同的形状。广播后的张量与输入张量共享相同的存储（即数据没有被复制）。
#
# 广播规则
# 从右到左对齐维度：
#
# 比较两个张量的维度，从后向前逐一对比。
# 如果维度数量不同，则较小的张量在缺失的维度上补充为大小为
# 1
# 的维度。
# 维度匹配规则：
#
# 如果两个张量在某个维度上的大小相同，则该维度可以兼容。
# 如果其中一个张量在某个维度的大小为
# 1，则该张量在此维度上会扩展为另一个张量的大小。
# 如果两个张量在某个维度的大小既不相等又不为
# 1，则广播失败，抛出错误。
# 结果维度：
#
# 广播后张量的结果维度是逐元素比较后的最大形状。


# 定义两个张量
a = torch.tensor([1, 2, 3])  # 形状为 (3,)
b = torch.tensor([[5], [6], [7], [8]])  # 形状为 (4, 1)

# 广播张量
# a的广播步骤：
#   1. a维度数量是1，b维度数量是2，首先将a的维度数量空缺个数从左到右，依次补充1，即 (3) -> (1, 3)
#   2. a和b的维度数量相等后，将a中dim==1的地方改为b对应位置的dim值，即 (1, 3) -> (4, 3)
# b的广播步骤：
#   1. b的维度数量和a已经相等，只需要把b中dim==1的地方改为a对应位置的dim值即可，即 (4, 1) -> (4, 3)
broadcasted_a, broadcasted_b = torch.broadcast_tensors(a, b)

print(broadcasted_a)
# tensor([[1, 2, 3],
#         [1, 2, 3],
#         [1, 2, 3],
#         [1, 2, 3]])

print(broadcasted_b)
# tensor([[5, 5, 5],
#         [6, 6, 6],
#         [7, 7, 7],
#         [8, 8, 8]])